import requests
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from sklearn.preprocessing import MinMaxScaler
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Bidirectional
import tushare as ts
import math
import random

# 设置随机种子，保证结果的可复现性
random.seed(1)
np.random.seed(1)
tf.random.set_seed(1)

# 设置中文字体为黑体，解决负号显示问题（用于部分和matplotlib相关的操作，比如后续模型训练结果可视化等）
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 获取股票数据的函数，使用Tushare的数据接口（需替换token才能正常获取数据）
def get_stock_data(token, ts_code, start_date, end_date, fields):
    """
    从tushare获取指定股票代码、日期范围以及字段的股票数据

    参数:
    token (str): 在tushare平台注册获取的token
    ts_code (str): 股票代码，如 "601939.SH"
    start_date (str): 开始日期，格式为 "YYYYMMDD"
    end_date (str): 结束日期，格式为 "YYYYMMDD"
    fields (list): 要获取的股票数据字段列表，如 ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    返回:
    pd.DataFrame: 获取到的股票数据，如果获取失败返回None
    """
    try:
        pro = ts.pro_api('572b6125c826e36619a670a9943995412179dabcfb739c5c48913f46')
        data = pro.daily(ts_code=ts_code,
                         start_date=start_date,
                         end_date=end_date,
                         fields=fields)
        return data
    except Exception as e:
        print(f"获取股票数据时出现错误: {e}")
        return None


# 数据预处理函数，进行缺失值处理、数据归一化等操作，增加更多详细信息打印
def preprocess_data(data):
    df = data.copy()
    # 将交易日期列设置为索引，并转换为日期时间类型
    df.set_index('trade_date', inplace=True)
    df.index = pd.to_datetime(df.index)
    print("设置日期索引后的基本信息：")
    print(df.info())
    print("设置日期索引后的列信息：", df.columns)
    print("设置日期索引后的数据形状:", df.shape)

    # 去除包含缺失值的行
    df.dropna(inplace=True)
    print("去除缺失值后的基本信息：")
    print(df.info())
    print("去除缺失值后的列信息：", df.columns)
    print("去除缺失值后的数据形状:", df.shape)

    # 提取需要的列（开盘价、最高价、最低价、收盘价、成交量），可根据实际需求调整
    columns = ['open', 'high', 'low', 'close', 'vol']
    df = df[columns]
    print("提取关键列后的基本信息：")
    print(df.info())
    print("提取关键列后的列信息：", df.columns)
    print("提取关键列后的数据形状:", df.shape)

    # 数据归一化，这里使用MinMaxScaler将数据归一化到0-1区间
    scaler = MinMaxScaler(feature_range=(0, 1))
    df_scaled = scaler.fit_transform(df)
    print("归一化后的数据形状:", df_scaled.shape)
    return df, df_scaled, scaler


# 创建训练数据集和测试数据集的函数，按照时间序列划分数据
def create_dataset(dataset, time_steps=1):
    dataX, dataY = [], []
    for i in range(len(dataset) - time_steps - 1):
        a = dataset[i:(i + time_steps), :]
        dataX.append(a)
        dataY.append(dataset[i + time_steps, -1])
    return np.array(dataX), np.array(dataY)


# 构建并训练LSTM模型的函数，支持不同模型类型选择
def build_and_train_model(model_type, X_train, y_train, X_test, y_test):
    model = Sequential()
    if model_type == 1:
        # 单层 LSTM
        model.add(LSTM(units=50, activation='relu',
                       input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(Dense(units=1))
    elif model_type == 2:
        # 多层 LSTM
        model.add(LSTM(units=50, activation='relu', return_sequences=True,
                       input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(LSTM(units=50, activation='relu'))
        model.add(Dense(1))
    elif model_type == 3:
        # 双向 LSTM
        model.add(Bidirectional(LSTM(50, activation='relu'),
                                input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(Dense(1))
    model.compile(loss='mean_squared_error', optimizer='adam')
    history = model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=30, batch_size=64, verbose=1)
    return model, history


# 计算均方根误差（RMSE），用于评估模型性能
def calculate_rmse(y_true, y_pred):
    return math.sqrt(((y_true - y_pred) ** 2).mean())


if __name__ == "__main__":
    # Tushare平台获取的token，需替换为你自己的真实token
    token = 'your_token_here'
    ts_code = "601939.SH"  # 建设银行股票代码，可替换为其他感兴趣的股票代码
    start_date = "20100101"
    end_date = "20241225"
    fields = ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    # 获取股票数据
    stock_data = get_stock_data(token, ts_code, start_date, end_date, fields)
    if stock_data is None:
        print("无法获取股票数据，程序退出。")
        exit(1)

    # 数据预处理
    df, df_scaled, scaler = preprocess_data(stock_data)

    # 划分训练集和测试集，这里简单按照80%作为训练集，20%作为测试集，你可以根据实际情况调整划分比例哦
    print("数据总长度:", len(df_scaled))
    train_size = int(len(df_scaled) * 0.8)
    print("计算出的训练集预期大小:", train_size)
    test_size = len(df_scaled) - train_size
    print("计算出的测试集预期大小:", test_size)
    train_data, test_data = df_scaled[0:train_size, :], df_scaled[train_size:len(df_scaled), :]
    print("实际划分后训练集数据形状:", train_data.shape)
    print("实际划分后测试集数据形状:", test_data.shape)
    time_steps = 10  # 时间步长，可根据需要调整
    X_train, y_train = create_dataset(train_data, time_steps)
    X_test, y_test = create_dataset(test_data, time_steps)
    # 调整输入数据形状，使其符合LSTM模型输入要求
    X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], df.shape[1]))
    X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], df.shape[1]))

    # 选择模型类型（1: 单层 LSTM; 2: 多层 LSTM; 3: 双向 LSTM）
    model_type = 3
    # 构建并训练模型
    model, history = build_and_train_model(model_type, X_train, y_train, X_test, y_test)

    # 绘制训练过程中的损失曲线，查看模型收敛情况（使用matplotlib）
    plt.plot(history.history['loss'], label='训练损失')
    plt.plot(history.history['val_loss'], label='验证损失')
    plt.title('模型训练损失曲线')
    plt.xlabel('Epoch')
    plt.ylabel('损失值')
    plt.legend()
    plt.show()

    # 使用训练好的模型进行预测
    train_predict = model.predict(X_train)

    # 调整 train_predict 的维度，使其与 train_data 匥配
    train_predict = np.reshape(train_predict, (train_predict.shape[0], 1))

    # 计算 train_data_trimmed 时，去除前 `time_steps` 行，确保和 `train_predict` 长度一致
    train_data_trimmed = train_data[time_steps:len(train_predict) + time_steps, 0:df.shape[1] - 1]

    # 进行反归一化操作时，确保 train_data 和 train_predict 在合并后的维度匹配
    train_predict = scaler.inverse_transform(np.concatenate((train_data_trimmed, train_predict), axis=1))[:, -1]


   

    print("train_data_trimmed shape:", train_data_trimmed.shape)
    print("train_predict shape:", train_predict.shape)


    # 计算训练集和测试集的均方根误差（RMSE），评估模型性能
    # 修剪 train_predict 的长度，确保与 y_true（即真实股价数据）长度一致
    train_predict_trimmed = train_predict[:len(df['close'].iloc[time_steps:len(train_predict) + time_steps - 1])]

   # 使用模型对测试集进行预测
    test_predict = model.predict(X_test)

# 调整 test_predict 的维度
    test_predict = np.reshape(test_predict, (test_predict.shape[0], 1))

# 修剪 test_predict 的长度，确保与 y_true 长度一致
    test_data_trimmed = test_data[time_steps:len(test_predict) + time_steps, 0:df.shape[1] - 1]

# 进行反归一化操作
    test_predict = scaler.inverse_transform(np.concatenate((test_data_trimmed, test_predict), axis=1))[:, -1]

# 计算测试集 RMSE
    test_rmse = calculate_rmse(df['close'].iloc[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1].values, test_predict)

# 输出测试集 RMSE
    print("测试集均方根误差（RMSE）:", test_rmse)

    test_rmse = calculate_rmse(df['close'].iloc[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1].values, test_predict)
    print("训练集均方根误差（RMSE）:", train_rmse)
    print("测试集均方根误差（RMSE）:", test_rmse)

    # 使用Plotly库绘制股票价格走势以及预测结果对比图
    fig = go.Figure()

    # 添加真实股价走势（收盘价）
    fig.add_trace(go.Scatter(x=df.index[time_steps:len(train_predict) + time_steps - 1],
                             y=df['close'].iloc[time_steps:len(train_predict) + time_steps - 1],
                             mode='lines',
                             name='真实股价（训练集）'))
    fig.add_trace(go.Scatter(x=df.index[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1],
                             y=df['close'].iloc[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1],
                             mode='lines',
                             name='真实股价（测试集）'))

    # 添加预测股价走势
    fig.add_trace(go.Scatter(x=df.index[time_steps:len(train_predict) + time_steps - 1],
                             y=train_predict,
                             mode='lines',
                             name='预测股价（训练集）',
                             line=dict(dash='dash')))
    fig.add_trace(go.Scatter(x=df.index[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1],
                             y=test_predict,
                             mode='lines',
                             name='预测股价（测试集）',
                             line=dict(dash='dash')))

    # 设置图表标题等布局信息
    fig.update_layout(title='股票真实价格与预测价格对比',
                      xaxis_title='日期',
                      yaxis_title='股价',
                      legend=dict(x=0.8, y=0.9),  # 调整图例位置，可根据喜好修改
                      hovermode='x unified')  # 设置悬停交互模式，方便查看具体数据点

    # 显示图表（如果是在Jupyter Notebook中会直接内联显示，在脚本中运行会在浏览器中打开图表页面）
    fig.show()

    # 额外添加一些之前可能有的中间数据查看或其他相关操作示例（以下只是简单示意，你可以根据实际之前代码情况调整哦）
    print("训练集预测结果的前5个数据：", train_predict[:5])
    print("测试集预测结果的前5个数据：", test_predict[:5])
    print("原始数据最后5行信息：")
    print(df.tail(5))